#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <assert.h>

#include <pthread.h>

pthread_rwlock_t rwlock;

#define READLOCK 1
#define WRITELOCK 2

void *pth_func(void *arg)
{
    int flag = (int)arg;

    if (READLOCK == flag)
    {
        if (pthread_rwlock_rdlock(&rwlock) == 0)
        {
            printf("Thread(%lx) Read Lock Success!\n", pthread_self());
        }
        else
        {
            printf("Thread(%lx) Read Lock Failure!\n", pthread_self());
        }
    }
    else if (WRITELOCK == flag)
    {
        if (pthread_rwlock_wrlock(&rwlock) == 0)
        {
            printf("Wirte Lock Success!\n");
        }
        else
        {
            printf("Wirte Lock Failure!\n");
        }
    }
    else
    {
        printf("flag error!\n");
        return (void *)-1;
    }

    printf("pthread(%lx) sleep...\n", pthread_self());
    sleep(5);

    printf("Lock(%s) pthread(%lx) leave!\n", flag == READLOCK ? "READLOCK" : "WRITELOCK", pthread_self());

    pthread_rwlock_unlock(&rwlock);
    return (void *)0;
}

int get_lock_flag(char *str)
{
    assert(str != NULL);

    if (strncmp("r", str, strlen(str)) == 0)
    {
        return READLOCK;
    }
    else if (strncmp("w", str, strlen(str)) == 0)
    {
        return WRITELOCK;
    }

    return 0;
}

int main(int argc, char **argv)
{

    if (argc < 4)
    {
        printf("Usage: %s <[r|w] [r|w] [r|w]>", argv[0]);
        exit(EXIT_FAILURE);
    }

    pthread_rwlock_init(&rwlock, NULL);

    pthread_t pth1, pth2, pth3;
    int pth1_arg, pth2_arg, pth3_arg;
    pth1_arg = get_lock_flag(argv[1]);
    pth2_arg = get_lock_flag(argv[2]);
    pth3_arg = get_lock_flag(argv[3]);

    if (pthread_create(&pth1, NULL, pth_func, (void *)pth1_arg) != 0)
    {
        printf("pthread_create failure!\n");
        exit(EXIT_FAILURE);
    }
    else
    {
        printf("pthread1(%lx) create success!\n", pth1);
    }

    sleep(1);
    if (pthread_create(&pth2, NULL, pth_func, (void *)pth2_arg) != 0)
    {
        printf("pthread_create failure!\n");
        exit(EXIT_FAILURE);
    }
    else
    {
        printf("pthread2(%lx) create success!\n", pth2);
    }
    sleep(2);
    if (pthread_create(&pth3, NULL, pth_func, (void *)pth3_arg) != 0)
    {
        printf("pthread_create failure!\n");
        exit(EXIT_FAILURE);
    }
    else
    {
        printf("pthread3(%lx) create success!\n", pth3);
    }

    pthread_join(pth1, NULL);
    pthread_join(pth2, NULL);
    pthread_join(pth3, NULL);

    printf("pthread1(%lx) lock[%s]\tpthread2(%lx) lock[%s]\tpthread3(%lx) lock[%s]\n",
           pth1, argv[1], pth2, argv[2], pth3, argv[3]);

    return 0;
}